/*
 * Copyright (c) 2010 by Aleksey Zhukov
 * Licensed under the Apache License, Version 2.0 (the "License")
 */

package ru.listfunc4j;

import com.sun.tools.attach.VirtualMachine;
import javassist.ClassPool;
import javassist.CtClass;
import sun.management.VMManagement;

import java.io.File;
import java.io.FileOutputStream;
import java.lang.instrument.ClassDefinition;
import java.lang.management.ManagementFactory;
import java.lang.management.RuntimeMXBean;
import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.util.jar.Attributes;
import java.util.jar.JarEntry;
import java.util.jar.JarOutputStream;
import java.util.jar.Manifest;

/**
 * User: J0k3r
 * Date: 24.10.2010
 * Time: 12:40:41
 */
public class AgentWork {
    static void init() {
        try {
            VirtualMachine machine = VirtualMachine.attach(getPidFromRuntimeMBean());
            final File jarFile = File.createTempFile("agent", ".jar");
            jarFile.deleteOnExit();

            final Manifest manifest = new Manifest();
            final Attributes mainAttributes = manifest.getMainAttributes();
            mainAttributes.put(Attributes.Name.MANIFEST_VERSION, "1.0");
            mainAttributes.put(new Attributes.Name("Agent-Class"), Agent.class.getName());
            mainAttributes.put(new Attributes.Name("Can-Retransform-Classes"), "true");
            mainAttributes.put(new Attributes.Name("Can-Redefine-Classes"), "true");

            final JarOutputStream jos = new JarOutputStream(new FileOutputStream(jarFile), manifest);
            final JarEntry agent = new JarEntry(Agent.class.getName().replace('.', '/') + ".class");
            jos.putNextEntry(agent);
            final ClassPool pool = ClassPool.getDefault();
            final CtClass ctClass = pool.get(Agent.class.getName());
            jos.write(ctClass.toBytecode());
            jos.closeEntry();
            jos.close();
            machine.loadAgent(jarFile.getAbsolutePath());
        } catch (Exception ex) {
            throw new RuntimeException(ex);
        }
    }

    public static void redefineClass(Class clazz, byte[] bytes) throws Exception {
        Agent.redefineClasses(new ClassDefinition(clazz, bytes));
    }

    private static String getPidFromRuntimeMBean() throws Exception {

        RuntimeMXBean mxbean = ManagementFactory.getRuntimeMXBean();
        Field jvmField = mxbean.getClass().getDeclaredField("jvm");

        jvmField.setAccessible(true);
        VMManagement management = (VMManagement) jvmField.get(mxbean);
        Method method = management.getClass().getDeclaredMethod("getProcessId");
        method.setAccessible(true);
        Integer processId = (Integer) method.invoke(management);

        return processId.toString();
    }
}
